package tornadoj.data;

import com.google.gson.Gson;
import org.apache.commons.lang3.StringUtils;
import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.analysis.standard.StandardAnalyzer;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
import org.apache.lucene.document.StoredField;
import org.apache.lucene.document.TextField;
import org.apache.lucene.index.*;
import org.apache.lucene.queryparser.classic.MultiFieldQueryParser;
import org.apache.lucene.queryparser.classic.QueryParser;
import org.apache.lucene.search.*;
import org.apache.lucene.search.highlight.*;
import org.apache.lucene.store.Directory;
import org.apache.lucene.store.FSDirectory;
import org.apache.lucene.util.Version;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import tornadoj.util.Utils;

import java.io.IOException;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 搜索工具类
 * Created by tinyms on 13-12-18.
 */
public class LuceneUtil {
    private final static Logger logger = LoggerFactory.getLogger(LuceneUtil.class);

    private Directory directory = null;
    public final static Version CURRENT = Version.LATEST;

    public LuceneUtil(String storedPath) {
        try {
            Utils.mkdirs(storedPath);
            directory = FSDirectory.open(Paths.get(storedPath));
        } catch (IOException e) {
            logger.info(e.getLocalizedMessage(), e);
        }
    }

    /**
     * JSON格式
     * [{"action":"add", "id":"1", "category":"PosturePicture", "meta":"display", "indexText":"Every Body, Good by!"}]
     *
     * @param json 索引数据
     */
    public void parseJSON(String json){
        java.lang.reflect.Type type = new com.google.gson.reflect.TypeToken<List<DocItem>>() {}.getType();
        Gson gson = new Gson();
        final List<DocItem> arr = gson.fromJson(json, type);
        this.parseDocItems(arr);
    }

    public void parseDocItems(List<DocItem> arr){
        if (arr==null){
            return;
        }
        this.write(writer -> {
            for (DocItem docItem: arr){
                String act = docItem.getAction();
                if ("add".equalsIgnoreCase(act)){
                    Document doc = new Document();
                    doc.add(new StoredField("id", docItem.getId()));
                    doc.add(new TextField("category", docItem.getCategory(), Field.Store.YES));
                    doc.add(new StoredField("meta", docItem.getMeta()));
                    doc.add(new TextField("indexText", docItem.getIndexText(), Field.Store.YES));
                    try {
                        writer.addDocument(doc);
                    } catch (IOException e) {
                        logger.info(e.getLocalizedMessage(), e);
                    }
                }else if ("update".equalsIgnoreCase(act)){
                    Document doc = new Document();
                    doc.add(new StoredField("id", docItem.getId()));
                    doc.add(new TextField("category", docItem.getCategory(), Field.Store.YES));
                    doc.add(new StoredField("meta", docItem.getMeta()));
                    doc.add(new TextField("indexText", docItem.getIndexText(), Field.Store.YES));
                    try {
                        writer.updateDocument(new Term("id", docItem.getId()), doc);
                    } catch (IOException e) {
                        logger.info(e.getLocalizedMessage(), e);
                    }
                }else if ("delete".equalsIgnoreCase(act)){
                    try {
                        writer.deleteDocuments(new Term("id", docItem.getId()));
                        writer.forceMergeDeletes();
                    } catch (IOException e) {
                        logger.info(e.getLocalizedMessage(), e);
                    }
                }
            }
        });
    }

    public Map<String, Object> search(String category, String kw, int currentPage, int pageSize, boolean color){
        Map<String, Object> result = this.paginate(category, kw, new String[]{"category", "indexText"}, currentPage, pageSize, new ILuceneSearch() {
            @Override
            public Object doInSearch(Document doc, Analyzer analyzer, Highlighter highlighter) {
                Map<String, String> item = new HashMap<String, String>();
                String id = doc.get("id");
                String meta = doc.get("meta");
                item.put("id", id);
                item.put("meta", meta);
                if (color){
                    String indexText = LuceneUtil.colorKeyWords(doc, analyzer, highlighter, "indexText");
                    item.put("summary", indexText);
                }
                return item;
            }
        });
        return result;
    }

    public void write(ILuceneWriter custom) {
        try {
            IndexWriterConfig config = new IndexWriterConfig(
                    new StandardAnalyzer(StandardAnalyzer.STOP_WORDS_SET));
            IndexWriter writer = new IndexWriter(directory, config);
            custom.doInWrite(writer);
            writer.commit();
            writer.close();
        } catch (IOException e) {
            logger.info(e.getLocalizedMessage(), e);
        }
    }

    public Map<String, Object> paginate(String category, String keyWord, String[] fieldNames, int curpage, int pageSize, ILuceneSearch iSearchResultHandler) {
        Map<String, Object> result = new HashMap<String, Object>();
        result.put("rows", 0);
        result.put("pages", 0);
        result.put("items", new ArrayList());
        keyWord = StringUtils.trimToEmpty(keyWord);
        if (StringUtils.isBlank(keyWord)) {
            return result;
        }
        if (!keyWord.contains("\"")) {
            List<String> chinese = Utils.parseChinese(keyWord);
            for (String str : chinese) {
                keyWord = StringUtils.replace(keyWord, str, "\"" + str + "\"");
            }
        }
        keyWord = category + " " + keyWord;
        try {
            if (curpage <= 0) {
                curpage = 1;
            }
            if (pageSize <= 0) {
                pageSize = 20;
            }

            int start = (curpage - 1) * pageSize;
            IndexReader reader = DirectoryReader.open(this.directory);
            IndexSearcher searcher = new IndexSearcher(reader);
            Analyzer analyzer = new StandardAnalyzer(StandardAnalyzer.STOP_WORDS_SET);

            QueryParser queryParser = new MultiFieldQueryParser(fieldNames, analyzer);
            queryParser.setDefaultOperator(QueryParser.AND_OPERATOR);
            Query query = queryParser.parse(keyWord);
            int hm = start + pageSize;
            TopScoreDocCollector res = TopScoreDocCollector.create(hm);
            searcher.search(query, res);

            Highlighter highlighter = createHighlighter(query);

            List<Object> items = new ArrayList<Object>();
            int rowCount = res.getTotalHits();
            int pages = (rowCount - 1) / pageSize + 1; //计算总页数
            result.put("rows", rowCount);
            result.put("pages", pages);
            TopDocs tds = res.topDocs(start, pageSize);
            ScoreDoc[] sd = tds.scoreDocs;
            for (ScoreDoc aSd : sd) {
                Document doc = reader.document(aSd.doc);
                Object item = iSearchResultHandler.doInSearch(doc, analyzer, highlighter);
                items.add(item);
            }
            result.put("items", items);
            return result;

        } catch (Exception e) {
            logger.info(e.getLocalizedMessage(), e);
        }

        return result;
    }

    public static void main(String[] args) {
        LuceneUtil lucene = new LuceneUtil("/Users/apple/sb");
        //lucene.parseJSON("[{\"action\":\"add\", \"id\":\"1\", \"category\":\"PosturePicture\", \"meta\":\"display\", \"indexText\":\"Every Body, Good by!\"}]");
        String r = Utils.encode(lucene.search("PosturePicture", "Good", 1, 20, true));
        System.out.println(r);
    }

    public static String getField(Document doc, String fieldName) {
        String val = doc.get(fieldName);
        if (StringUtils.isBlank(val)) {
            return "";
        }
        return val;
    }

    public static String colorKeyWords(Document doc, Analyzer analyzer, Highlighter highlighter, String fieldName) {
        String val = doc.get(fieldName);
        if (StringUtils.isBlank(val)) {
            return "";
        }
        try {
            String colorText = highlighter.getBestFragment(analyzer, fieldName, val);
            if (!StringUtils.isBlank(colorText)) {
                return colorText;
            }
        } catch (IOException | InvalidTokenOffsetsException e) {
            logger.info(e.getLocalizedMessage(), e);
        }
        return val;
    }

    private static Highlighter createHighlighter(Query query) {
        SimpleHTMLFormatter simpleHTMLFormatter = new SimpleHTMLFormatter("<span style='color:red'>", "</span>");
        Highlighter highlighter = new Highlighter(simpleHTMLFormatter, new QueryScorer(query));
        SimpleFragmenter sf = new SimpleFragmenter(100);
        highlighter.setTextFragmenter(sf);
        return highlighter;
    }

    /**
     * 删除索引 删除的索引会保存到一个新的文件中（以del为结尾的文件 相当于删除到回收站）
     *
     * @param force 是否强制清空
     * @throws java.io.IOException
     */
    public void delete(String id, boolean force) throws IOException {
        IndexWriterConfig config = new IndexWriterConfig(
                new StandardAnalyzer(StandardAnalyzer.STOP_WORDS_SET));
        IndexWriter writer = new IndexWriter(directory, config);
        writer.deleteDocuments(new Term("id", id));
        if (force) {
            writer.forceMergeDeletes();
        }
        writer.close();
    }

    /**
     * 删除索引 删除的索引会保存到一个新的文件中（以del为结尾的文件 相当于删除到回收站）
     *
     * @param fieldName  字段名
     * @param fieldValue 字段值
     * @param force      是否强制清空
     * @throws java.io.IOException
     */
    public void delete(String fieldName, String fieldValue, boolean force) throws IOException {
        IndexWriterConfig config = new IndexWriterConfig(
                new StandardAnalyzer(StandardAnalyzer.STOP_WORDS_SET));
        IndexWriter writer = new IndexWriter(directory, config);
        writer.deleteDocuments(new Term(fieldName, fieldValue));
        if (force) {
            writer.forceMergeDeletes();
        }
        writer.close();
    }

    /**
     * 删除索引 删除的索引会保存到一个新的文件中（以del为结尾的文件 相当于删除到回收站）
     *
     * @throws java.io.IOException
     */
    public void update(String id, Document doc) throws IOException {
        IndexWriterConfig config = new IndexWriterConfig(
                new StandardAnalyzer(StandardAnalyzer.STOP_WORDS_SET));
        IndexWriter writer = new IndexWriter(directory, config);
        writer.updateDocument(new Term("id", id), doc);
        writer.commit();
        writer.close();
    }

    /**
     * 删除索引 删除的索引会保存到一个新的文件中（以del为结尾的文件 相当于删除到回收站）
     *
     * @throws java.io.IOException
     */
    public void update(String fieldName, String fieldValue, Document doc) throws IOException {
        IndexWriterConfig config = new IndexWriterConfig(
                new StandardAnalyzer(StandardAnalyzer.STOP_WORDS_SET));
        IndexWriter writer = new IndexWriter(directory, config);
        writer.updateDocument(new Term(fieldName, fieldValue), doc);
        writer.commit();
        writer.close();
    }

    /**
     * 删除所有的索引 删除的索引会保存到一个新的文件中（以del为结尾的文件 相当于删除到回收站）
     *
     * @throws java.io.IOException
     */
    public void deleteAll() throws IOException {
        IndexWriterConfig config = new IndexWriterConfig(
                new StandardAnalyzer(StandardAnalyzer.STOP_WORDS_SET));
        IndexWriter writer = new IndexWriter(directory, config);
        writer.deleteAll();
        writer.close();
    }

    /**
     * 删除已经删除的索引 对应上一个删除方法 清空回收站的文件
     *
     * @throws java.io.IOException
     */
    public void forceMergeDeletes() throws IOException {
        IndexWriterConfig config = new IndexWriterConfig(
                new StandardAnalyzer(StandardAnalyzer.STOP_WORDS_SET));
        IndexWriter writer = new IndexWriter(directory, config);
        writer.forceMergeDeletes();
        writer.close();
    }
}
